{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wget https://raw.githubusercontent.com/UniversalDependencies/UD_English-EWT/master/en_ewt-ud-dev.conllu\n",
    "# !wget https://raw.githubusercontent.com/UniversalDependencies/UD_English-EWT/master/en_ewt-ud-train.conllu\n",
    "# !wget https://raw.githubusercontent.com/UniversalDependencies/UD_English-EWT/master/en_ewt-ud-test.conllu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/husein/.local/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:516: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/home/husein/.local/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:517: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/home/husein/.local/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:518: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/home/husein/.local/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:519: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/home/husein/.local/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:520: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/home/husein/.local/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:525: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n",
      "/home/husein/.local/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:541: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/home/husein/.local/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:542: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/home/husein/.local/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:543: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/home/husein/.local/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:544: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/home/husein/.local/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:545: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/home/husein/.local/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:550: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n"
     ]
    }
   ],
   "source": [
    "import malaya\n",
    "import re\n",
    "from malaya.texts._text_functions import split_into_sentences\n",
    "from malaya.texts import _regex\n",
    "import numpy as np\n",
    "import itertools\n",
    "\n",
    "tokenizer = malaya.preprocessing._tokenizer\n",
    "splitter = split_into_sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_number_regex(s):\n",
    "    if re.match(\"^\\d+?\\.\\d+?$\", s) is None:\n",
    "        return s.isdigit()\n",
    "    return True\n",
    "\n",
    "def preprocessing(w):\n",
    "    if is_number_regex(w):\n",
    "        return '<NUM>'\n",
    "    elif re.match(_regex._money, w):\n",
    "        return '<MONEY>'\n",
    "    elif re.match(_regex._date, w):\n",
    "        return '<DATE>'\n",
    "    elif re.match(_regex._expressions['email'], w):\n",
    "        return '<EMAIL>'\n",
    "    elif re.match(_regex._expressions['url'], w):\n",
    "        return '<URL>'\n",
    "    else:\n",
    "        w = ''.join(''.join(s)[:2] for _, s in itertools.groupby(w))\n",
    "        return w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "({'PAD': 0,\n",
       "  'UNK': 1,\n",
       "  '_ROOT': 2,\n",
       "  '<NUM>': 3,\n",
       "  '<MONEY>': 4,\n",
       "  '<DATE>': 5,\n",
       "  '<URL>': 6,\n",
       "  '<EMAIL>': 7},\n",
       " {'PAD': 0,\n",
       "  'UNK': 1,\n",
       "  '_ROOT': 2,\n",
       "  '<NUM>': 3,\n",
       "  '<MONEY>': 4,\n",
       "  '<DATE>': 5,\n",
       "  '<URL>': 6,\n",
       "  '<EMAIL>': 7})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "word2idx = {'PAD': 0,'UNK':1, '_ROOT': 2}\n",
    "tag2idx = {'PAD': 0, '_<ROOT>': 1}\n",
    "char2idx = {'PAD': 0,'UNK':1, '_ROOT': 2}\n",
    "word_idx = 3\n",
    "tag_idx = 2\n",
    "char_idx = 3\n",
    "\n",
    "special_tokens = ['<NUM>', '<MONEY>', '<DATE>', '<URL>', '<EMAIL>']\n",
    "\n",
    "for t in special_tokens:\n",
    "    word2idx[t] = word_idx\n",
    "    word_idx += 1\n",
    "    char2idx[t] = char_idx\n",
    "    char_idx += 1\n",
    "    \n",
    "word2idx, char2idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "PAD = \"_PAD\"\n",
    "PAD_POS = \"_PAD_POS\"\n",
    "PAD_TYPE = \"_<PAD>\"\n",
    "PAD_CHAR = \"_PAD_CHAR\"\n",
    "ROOT = \"_ROOT\"\n",
    "ROOT_POS = \"_ROOT_POS\"\n",
    "ROOT_TYPE = \"_<ROOT>\"\n",
    "ROOT_CHAR = \"_ROOT_CHAR\"\n",
    "END = \"_END\"\n",
    "END_POS = \"_END_POS\"\n",
    "END_TYPE = \"_<END>\"\n",
    "END_CHAR = \"_END_CHAR\"\n",
    "\n",
    "def process_corpus(corpus, until = None):\n",
    "    global word2idx, tag2idx, char2idx, word_idx, tag_idx, char_idx\n",
    "    sentences, words, depends, labels, pos, chars = [], [], [], [], [], []\n",
    "    temp_sentence, temp_word, temp_depend, temp_label, temp_pos = [], [], [], [], []\n",
    "    first_time = True\n",
    "    for sentence in corpus:\n",
    "        try:\n",
    "            if len(sentence):\n",
    "                if sentence[0] == '#':\n",
    "                    continue\n",
    "                if first_time:\n",
    "                    print(sentence)\n",
    "                    first_time = False\n",
    "                sentence = sentence.split('\\t')\n",
    "                for c in sentence[1]:\n",
    "                    if c not in char2idx:\n",
    "                        char2idx[c] = char_idx\n",
    "                        char_idx += 1\n",
    "                if sentence[7] not in tag2idx:\n",
    "                    tag2idx[sentence[7]] = tag_idx\n",
    "                    tag_idx += 1\n",
    "                sentence[1] = preprocessing(sentence[1])\n",
    "                if sentence[1] not in word2idx:\n",
    "                    word2idx[sentence[1]] = word_idx\n",
    "                    word_idx += 1\n",
    "                temp_word.append(word2idx[sentence[1]])\n",
    "                temp_depend.append(int(sentence[6]))\n",
    "                temp_label.append(tag2idx[sentence[7]])\n",
    "                temp_sentence.append(sentence[1])\n",
    "                temp_pos.append(sentence[3])\n",
    "            else:\n",
    "                if len(temp_sentence) < 2 or len(temp_word) != len(temp_label):\n",
    "                    temp_word = []\n",
    "                    temp_depend = []\n",
    "                    temp_label = []\n",
    "                    temp_sentence = []\n",
    "                    temp_pos = []\n",
    "                    continue\n",
    "                words.append([word2idx['_ROOT']] + temp_word)\n",
    "                depends.append([0] + temp_depend)\n",
    "                labels.append([tag2idx['_<ROOT>']] + temp_label)\n",
    "                sentences.append([ROOT] + temp_sentence)\n",
    "                pos.append([ROOT_POS] + temp_pos)\n",
    "                char_ = [[char2idx['_ROOT']]]\n",
    "                for w in temp_sentence:\n",
    "                    if w in char2idx:\n",
    "                        char_.append([char2idx[w]])\n",
    "                    else:\n",
    "                        char_.append([char2idx[c] for c in w])\n",
    "                chars.append(char_)\n",
    "                temp_word = []\n",
    "                temp_depend = []\n",
    "                temp_label = []\n",
    "                temp_sentence = []\n",
    "                temp_pos = []\n",
    "        except Exception as e:\n",
    "            print(e, sentence)\n",
    "    return sentences[:-1], words[:-1], depends[:-1], labels[:-1], pos[:-1], chars[:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _obtain_child_index_for_left2right(heads):\n",
    "    child_ids = [[] for _ in range(len(heads))]\n",
    "    # skip the symbolic root.\n",
    "    for child in range(1, len(heads)):\n",
    "        head = heads[child]\n",
    "        child_ids[head].append(child)\n",
    "    return child_ids\n",
    "\n",
    "\n",
    "def _obtain_child_index_for_inside_out(heads):\n",
    "    child_ids = [[] for _ in range(len(heads))]\n",
    "    for head in range(len(heads)):\n",
    "        # first find left children inside-out\n",
    "        for child in reversed(range(1, head)):\n",
    "            if heads[child] == head:\n",
    "                child_ids[head].append(child)\n",
    "        # second find right children inside-out\n",
    "        for child in range(head + 1, len(heads)):\n",
    "            if heads[child] == head:\n",
    "                child_ids[head].append(child)\n",
    "    return child_ids\n",
    "\n",
    "\n",
    "def _obtain_child_index_for_depth(heads, reverse):\n",
    "    def calc_depth(head):\n",
    "        children = child_ids[head]\n",
    "        max_depth = 0\n",
    "        for child in children:\n",
    "            depth = calc_depth(child)\n",
    "            child_with_depth[head].append((child, depth))\n",
    "            max_depth = max(max_depth, depth + 1)\n",
    "        child_with_depth[head] = sorted(child_with_depth[head], key=lambda x: x[1], reverse=reverse)\n",
    "        return max_depth\n",
    "\n",
    "    child_ids = _obtain_child_index_for_left2right(heads)\n",
    "    child_with_depth = [[] for _ in range(len(heads))]\n",
    "    calc_depth(0)\n",
    "    return [[child for child, depth in child_with_depth[head]] for head in range(len(heads))]\n",
    "\n",
    "\n",
    "def _generate_stack_inputs(heads, types, prior_order):\n",
    "    if prior_order == 'deep_first':\n",
    "        child_ids = _obtain_child_index_for_depth(heads, True)\n",
    "    elif prior_order == 'shallow_first':\n",
    "        child_ids = _obtain_child_index_for_depth(heads, False)\n",
    "    elif prior_order == 'left2right':\n",
    "        child_ids = _obtain_child_index_for_left2right(heads)\n",
    "    elif prior_order == 'inside_out':\n",
    "        child_ids = _obtain_child_index_for_inside_out(heads)\n",
    "    else:\n",
    "        raise ValueError('Unknown prior order: %s' % prior_order)\n",
    "\n",
    "    stacked_heads = []\n",
    "    children = []\n",
    "    siblings = []\n",
    "    stacked_types = []\n",
    "    skip_connect = []\n",
    "    prev = [0 for _ in range(len(heads))]\n",
    "    sibs = [0 for _ in range(len(heads))]\n",
    "    stack = [0]\n",
    "    position = 1\n",
    "    while len(stack) > 0:\n",
    "        head = stack[-1]\n",
    "        stacked_heads.append(head)\n",
    "        siblings.append(sibs[head])\n",
    "        child_id = child_ids[head]\n",
    "        skip_connect.append(prev[head])\n",
    "        prev[head] = position\n",
    "        if len(child_id) == 0:\n",
    "            children.append(head)\n",
    "            sibs[head] = 0\n",
    "            stacked_types.append(tag2idx['PAD'])\n",
    "            stack.pop()\n",
    "        else:\n",
    "            child = child_id.pop(0)\n",
    "            children.append(child)\n",
    "            sibs[head] = child\n",
    "            stack.append(child)\n",
    "            stacked_types.append(types[child])\n",
    "        position += 1\n",
    "\n",
    "    return stacked_heads, children, siblings, stacked_types, skip_connect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\tFrom\tfrom\tADP\tIN\t_\t3\tcase\t3:case\t_\n",
      "invalid literal for int() with base 10: '_' ['10.1', 'has', 'have', 'VERB', 'VBZ', '_', '_', '_', '8:parataxis', 'CopyOf=-1']\n",
      "invalid literal for int() with base 10: '_' ['21.1', 'has', 'have', 'VERB', 'VBZ', '_', '_', '_', '16:conj:and', 'CopyOf=-1']\n"
     ]
    }
   ],
   "source": [
    "with open('en_ewt-ud-dev.conllu') as fopen:\n",
    "    dev = fopen.read().split('\\n')\n",
    "\n",
    "sentences_dev, words_dev, depends_dev, labels_dev, _, seq_dev = process_corpus(dev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "stacked_heads_test, children_test, siblings_test, stacked_types_test = [], [], [], []\n",
    "for i in range(len(sentences_dev)):\n",
    "    stacked_heads, children, siblings, stacked_types, _ = _generate_stack_inputs(depends_dev[i], \n",
    "                                                                                 labels_dev[i], 'deep_first')\n",
    "    stacked_heads_test.append(stacked_heads)\n",
    "    children_test.append(children)\n",
    "    siblings_test.append(siblings)\n",
    "    stacked_types_test.append(stacked_types)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\tWhat\twhat\tPRON\tWP\tPronType=Int\t0\troot\t0:root\t_\n",
      "invalid literal for int() with base 10: '_' ['24.1', 'left', 'left', 'VERB', 'VBN', 'Tense=Past|VerbForm=Part', '_', '_', '6:parataxis', 'CopyOf=6']\n"
     ]
    }
   ],
   "source": [
    "with open('en_ewt-ud-test.conllu') as fopen:\n",
    "    test = fopen.read().split('\\n')\n",
    "\n",
    "sentences_test, words_test, depends_test, labels_test, _, seq_test = process_corpus(test)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(sentences_test)):\n",
    "    stacked_heads, children, siblings, stacked_types, _ = _generate_stack_inputs(depends_test[i], \n",
    "                                                                                 labels_test[i], 'deep_first')\n",
    "    stacked_heads_test.append(stacked_heads)\n",
    "    children_test.append(children)\n",
    "    siblings_test.append(siblings)\n",
    "    stacked_types_test.append(stacked_types)\n",
    "    \n",
    "sentences_test.extend(sentences_dev)\n",
    "words_test.extend(words_dev)\n",
    "depends_test.extend(depends_dev)\n",
    "labels_test.extend(labels_dev)\n",
    "seq_test.extend(seq_dev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\tAl\tAl\tPROPN\tNNP\tNumber=Sing\t0\troot\t0:root\tSpaceAfter=No\n",
      "invalid literal for int() with base 10: '_' ['8.1', 'reported', 'report', 'VERB', 'VBN', 'Tense=Past|VerbForm=Part|Voice=Pass', '_', '_', '5:conj:and', 'CopyOf=5']\n",
      "invalid literal for int() with base 10: '_' ['22.1', 'used', 'use', 'VERB', 'VBN', 'Tense=Past|VerbForm=Part', '_', '_', '13:advcl:with|17:conj:and', 'CopyOf=17']\n",
      "invalid literal for int() with base 10: '_' ['22.1', 'used', 'use', 'VERB', 'VBN', 'Tense=Past|VerbForm=Part', '_', '_', '13:advcl:with|17:conj:and', 'CopyOf=17']\n",
      "invalid literal for int() with base 10: '_' ['11.1', 'called', 'call', 'VERB', 'VBN', 'Tense=Past|VerbForm=Part|Voice=Pass', '_', '_', '3:conj:and', 'CopyOf=3']\n",
      "invalid literal for int() with base 10: '_' ['14.1', 'is', 'be', 'VERB', 'VBZ', '_', '_', '_', '1:conj:and', 'CopyOf=1']\n",
      "invalid literal for int() with base 10: '_' ['20.1', 'reflect', 'reflect', 'VERB', 'VBP', 'Mood=Ind|Tense=Pres|VerbForm=Fin', '_', '_', '7:acl:relcl|9:conj', 'CopyOf=9']\n",
      "invalid literal for int() with base 10: '_' ['21.1', 'recruited', 'recruit', 'VERB', 'VBD', 'Mood=Ind|Tense=Past|VerbForm=Fin', '_', '_', '9:conj:and', 'CopyOf=9']\n",
      "invalid literal for int() with base 10: '_' ['9.1', 'wish', 'wish', 'VERB', 'VBP', 'Mood=Ind|Tense=Pres|VerbForm=Fin', '_', '_', '2:conj:and', 'CopyOf=2']\n",
      "invalid literal for int() with base 10: '_' ['38.1', 'supplied', 'supply', 'VERB', 'VBN', 'Tense=Past|VerbForm=Part|Voice=Pass', '_', '_', '16:conj:and', 'CopyOf=16']\n",
      "invalid literal for int() with base 10: '_' ['18.1', 'keep', 'keep', 'VERB', 'VB', 'Mood=Imp|VerbForm=Fin', '_', '_', '14:conj:and', 'CopyOf=14']\n",
      "invalid literal for int() with base 10: '_' ['21.1', 'keep', 'keep', 'VERB', 'VB', 'Mood=Imp|VerbForm=Fin', '_', '_', '14:conj:and', 'CopyOf=14']\n",
      "invalid literal for int() with base 10: '_' ['18.1', 'mean', 'mean', 'VERB', 'VB', 'VerbForm=Inf', '_', '_', '8:conj', 'CopyOf=8']\n",
      "invalid literal for int() with base 10: '_' ['30.1', 'play', 'play', 'VERB', 'VBP', 'Mood=Ind|Tense=Pres|VerbForm=Fin', '_', '_', '18:acl:relcl|27:conj:but', 'CopyOf=27']\n",
      "invalid literal for int() with base 10: '_' ['22.1', 'have', 'have', 'VERB', 'VBP', 'Mood=Ind|Tense=Pres|VerbForm=Fin', '_', '_', '17:conj', 'CopyOf=17']\n",
      "invalid literal for int() with base 10: '_' ['27.1', 'have', 'have', 'VERB', 'VBP', 'Mood=Ind|Tense=Pres|VerbForm=Fin', '_', '_', '17:conj', 'CopyOf=17']\n",
      "invalid literal for int() with base 10: '_' ['49.1', 'helped', 'help', 'VERB', 'VBD', '_', '_', '_', '38:conj:but', 'CopyOf=38']\n",
      "invalid literal for int() with base 10: '_' ['7.1', 'found', 'find', 'VERB', 'VBD', 'Mood=Ind|Tense=Past|VerbForm=Fin', '_', '_', '3:conj', 'CopyOf=3']\n",
      "invalid literal for int() with base 10: '_' ['10.1', 'excited', 'excited', 'ADJ', 'JJ', 'Degree=Pos', '_', '_', '4:advcl', 'CopyOf=4']\n",
      "invalid literal for int() with base 10: '_' ['15.1', \"'s\", 'be', 'VERB', 'VBZ', '_', '_', '_', '2:conj:and', 'CopyOf=2']\n",
      "invalid literal for int() with base 10: '_' ['25.1', 'took', 'take', 'VERB', 'VBD', 'Mood=Ind|Tense=Past|VerbForm=Fin', '_', '_', '17:conj:and', 'CopyOf=17']\n",
      "invalid literal for int() with base 10: '_' ['10.1', 'loss', 'lose', 'VERB', 'VBD', 'Mood=Ind|Tense=Past|VerbForm=Fin', '_', '_', '3:conj:and', 'CopyOf=3']\n",
      "invalid literal for int() with base 10: '_' ['11.1', 'leave', 'leave', 'VERB', 'VB', 'VerbForm=Inf', '_', '_', '7:parataxis', 'CopyOf=7']\n",
      "invalid literal for int() with base 10: '_' ['24.1', 'charge', 'charge', 'VERB', 'VBP', 'Mood=Ind|Tense=Pres|VerbForm=Fin', '_', '_', '16:conj:and', 'CopyOf=16']\n"
     ]
    }
   ],
   "source": [
    "with open('en_ewt-ud-train.conllu') as fopen:\n",
    "    train = fopen.read().split('\\n')\n",
    "\n",
    "sentences_train, words_train, depends_train, labels_train, _, _ = process_corpus(train)\n",
    "\n",
    "stacked_heads_train, children_train, siblings_train, stacked_types_train = [], [], [], []\n",
    "for i in range(len(sentences_train)):\n",
    "    stacked_heads, children, siblings, stacked_types, _ = _generate_stack_inputs(depends_train[i], \n",
    "                                                                                 labels_train[i], 'deep_first')\n",
    "    stacked_heads_train.append(stacked_heads)\n",
    "    children_train.append(children)\n",
    "    siblings_train.append(siblings)\n",
    "    stacked_types_train.append(stacked_types)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(12000, 3824)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(sentences_train), len(sentences_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "21974"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idx2word = {v:k for k, v in word2idx.items()}\n",
    "idx2tag = {v:k for k, v in tag2idx.items()}\n",
    "len(idx2word)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import Enum\n",
    "\n",
    "class PriorOrder(Enum):\n",
    "    DEPTH = 0\n",
    "    INSIDE_OUT = 1\n",
    "    LEFT2RIGTH = 2\n",
    "\n",
    "class BiAAttention:\n",
    "    def __init__(self, input_size_encoder, input_size_decoder, num_labels):\n",
    "        self.input_size_encoder = input_size_encoder\n",
    "        self.input_size_decoder = input_size_decoder\n",
    "        self.num_labels = num_labels\n",
    "        \n",
    "        self.W_d = tf.get_variable(\"W_d\", shape=[self.num_labels, self.input_size_decoder],\n",
    "           initializer=tf.contrib.layers.xavier_initializer())\n",
    "        self.W_e = tf.get_variable(\"W_e\", shape=[self.num_labels, self.input_size_encoder],\n",
    "           initializer=tf.contrib.layers.xavier_initializer())\n",
    "        self.U = tf.get_variable(\"U\", shape=[self.num_labels, self.input_size_decoder, self.input_size_encoder],\n",
    "           initializer=tf.contrib.layers.xavier_initializer())\n",
    "        \n",
    "    def forward(self, input_d, input_e, mask_d=None, mask_e=None):\n",
    "        batch = tf.shape(input_d)[0]\n",
    "        length_decoder = tf.shape(input_d)[1]\n",
    "        length_encoder = tf.shape(input_e)[1]\n",
    "        out_d = tf.expand_dims(tf.matmul(self.W_d, tf.transpose(input_d, [0, 2, 1])), 3)\n",
    "        out_e = tf.expand_dims(tf.matmul(self.W_e, tf.transpose(input_e, [0, 2, 1])), 2)\n",
    "        output = tf.matmul(tf.expand_dims(input_d, 1), self.U)\n",
    "        output = tf.matmul(output, tf.transpose(tf.expand_dims(input_e, 1), [0, 1, 3, 2]))\n",
    "        \n",
    "        output = output + out_d + out_e\n",
    "        \n",
    "        if mask_d is not None:\n",
    "            d = tf.expand_dims(tf.expand_dims(mask_d, 1), 3)\n",
    "            e = tf.expand_dims(tf.expand_dims(mask_e, 1), 2)\n",
    "            output = output * d * e\n",
    "            \n",
    "        return output\n",
    "    \n",
    "class BiLinear:\n",
    "    def __init__(self, left_features, right_features, out_features):\n",
    "        self.left_features = left_features\n",
    "        self.right_features = right_features\n",
    "        self.out_features = out_features\n",
    "        \n",
    "        self.U = tf.get_variable(\"U-bi\", shape=[out_features, left_features, right_features],\n",
    "           initializer=tf.contrib.layers.xavier_initializer())\n",
    "        self.W_l = tf.get_variable(\"Wl\", shape=[out_features, left_features],\n",
    "           initializer=tf.contrib.layers.xavier_initializer())\n",
    "        self.W_r = tf.get_variable(\"Wr\", shape=[out_features, right_features],\n",
    "           initializer=tf.contrib.layers.xavier_initializer())\n",
    "    \n",
    "    def forward(self, input_left, input_right):\n",
    "        left_size = tf.shape(input_left)\n",
    "        output_shape = tf.concat([left_size[:-1], [self.out_features]], axis = 0)\n",
    "        batch = tf.cast(tf.reduce_prod(left_size[:-1]), tf.int32)\n",
    "        input_left = tf.reshape(input_left, (batch, self.left_features))\n",
    "        input_right = tf.reshape(input_right, (batch, self.right_features))\n",
    "        tiled = tf.tile(tf.expand_dims(input_left, axis = 0), (self.out_features,1,1))\n",
    "        output = tf.transpose(tf.reduce_sum(tf.matmul(tiled, self.U), axis = 2))\n",
    "        output = output + tf.matmul(input_left, tf.transpose(self.W_l))\\\n",
    "        + tf.matmul(input_right, tf.transpose(self.W_r))\n",
    "        \n",
    "        return tf.reshape(output, output_shape)\n",
    "\n",
    "class StackPointer:\n",
    "    def __init__(self, word_dim, num_words, char_dim, num_chars, num_filters, kernel_size,\n",
    "                 input_size_decoder, hidden_size, layers,\n",
    "                 num_labels, arc_space, type_space):\n",
    "        \n",
    "        def cells(size, reuse=False):\n",
    "            return tf.nn.rnn_cell.LSTMCell(size,\n",
    "                                           initializer=tf.orthogonal_initializer(),reuse=reuse,\n",
    "                                           state_is_tuple=False)\n",
    "        \n",
    "        self.word_embedd = tf.Variable(tf.random_uniform([num_words, word_dim], -1, 1))\n",
    "        self.char_embedd = tf.Variable(tf.random_uniform([num_chars, char_dim], -1, 1))\n",
    "        self.conv1d = tf.layers.Conv1D(num_filters, kernel_size, 1, padding='VALID')\n",
    "        self.num_labels = num_labels\n",
    "        self.prior_order = PriorOrder.DEPTH\n",
    "        self.char_dim = char_dim\n",
    "        self.layers = layers\n",
    "        self.encoder = tf.nn.rnn_cell.MultiRNNCell([cells(hidden_size) for _ in range(layers)],\n",
    "                                                   state_is_tuple=False)\n",
    "        self.encoder_char = tf.nn.rnn_cell.MultiRNNCell([cells(hidden_size) for _ in range(layers)],\n",
    "                                                   state_is_tuple=False)\n",
    "        self.decoder = tf.nn.rnn_cell.MultiRNNCell([cells(hidden_size) for _ in range(layers)],\n",
    "                                                   state_is_tuple=False)\n",
    "        self.hidden_size = hidden_size\n",
    "        self.arc_space = arc_space\n",
    "        \n",
    "        \n",
    "        self.src_dense = tf.layers.Dense(hidden_size)\n",
    "        self.hx_dense = tf.layers.Dense(hidden_size)\n",
    "\n",
    "        self.arc_h = tf.layers.Dense(arc_space)\n",
    "        self.arc_c = tf.layers.Dense(arc_space)\n",
    "        self.attention = BiAAttention(arc_space, arc_space, 1)\n",
    "\n",
    "        self.type_h = tf.layers.Dense(type_space)\n",
    "        self.type_c = tf.layers.Dense(type_space)\n",
    "        self.bilinear = BiLinear(type_space, type_space, self.num_labels)\n",
    "        \n",
    "    def encode(self, input_word, input_char):\n",
    "        word = tf.nn.embedding_lookup(self.word_embedd, input_word)\n",
    "        char = tf.nn.embedding_lookup(self.char_embedd, input_char)\n",
    "        s = tf.shape(char)\n",
    "        char = tf.reshape(\n",
    "            char, shape = [s[0] * s[1], s[-2], self.char_dim]\n",
    "        )\n",
    "        output, _ = tf.nn.dynamic_rnn(self.encoder_char, char, dtype = tf.float32,\n",
    "                                      scope = 'encoder-char')\n",
    "        output = tf.reshape(\n",
    "            output[:, -1], shape = [s[0], s[1], self.hidden_size]\n",
    "        )\n",
    "        word_embedded = tf.concat([word, output], axis = -1)\n",
    "        output, hn = tf.nn.dynamic_rnn(self.encoder, word_embedded, dtype = tf.float32,\n",
    "                                      scope = 'encoder')\n",
    "        return output, hn\n",
    "    \n",
    "    def decode(self, output_encoder, heads, heads_stack, siblings, hn):\n",
    "        batch = tf.shape(output_encoder)[0]\n",
    "        batch_index = tf.range(0, batch)\n",
    "        t = tf.transpose(heads_stack)\n",
    "        broadcasted = tf.broadcast_to(batch_index, tf.shape(t))\n",
    "        concatenated = tf.transpose(tf.concat([tf.expand_dims(broadcasted, axis = 0), \n",
    "                                               tf.expand_dims(t, axis = 0)], axis = 0))\n",
    "        src_encoding = tf.gather_nd(output_encoder, concatenated)\n",
    "        \n",
    "        mask_sibs = tf.expand_dims(tf.cast(tf.not_equal(siblings, 0), tf.float32), axis = 2)\n",
    "        t = tf.transpose(siblings)\n",
    "        broadcasted = tf.broadcast_to(batch_index, tf.shape(t))\n",
    "        concatenated = tf.transpose(tf.concat([tf.expand_dims(broadcasted, axis = 0), \n",
    "                                               tf.expand_dims(t, axis = 0)], axis = 0))\n",
    "        output_enc_sibling = tf.gather_nd(output_encoder, concatenated) * mask_sibs\n",
    "        src_encoding = src_encoding + output_enc_sibling\n",
    "        \n",
    "        t = tf.transpose(heads_stack)\n",
    "        broadcasted = tf.broadcast_to(batch_index, tf.shape(t))\n",
    "        concatenated = tf.transpose(tf.concat([tf.expand_dims(broadcasted, axis = 0), \n",
    "                                       tf.expand_dims(t, axis = 0)],axis = 0))\n",
    "        g = tf.transpose(tf.gather_nd(heads, concatenated))\n",
    "        broadcasted = tf.broadcast_to(batch_index, tf.shape(g))\n",
    "        concatenated = tf.transpose(tf.concat([tf.expand_dims(broadcasted, axis = 0), \n",
    "                                               tf.expand_dims(g, axis = 0)],axis = 0))\n",
    "        output_enc_gpar = tf.gather_nd(output_encoder, concatenated)\n",
    "        src_encoding = src_encoding + output_enc_gpar\n",
    "        \n",
    "        src_encoding = tf.nn.elu(self.src_dense(src_encoding))\n",
    "        output, hn = tf.nn.dynamic_rnn(self.decoder, src_encoding, dtype = tf.float32,\n",
    "                                      initial_state = hn,\n",
    "                                      scope = 'decoder')\n",
    "        return output, hn\n",
    "    \n",
    "    def loss(self, input_word, input_char, \n",
    "             heads, stacked_heads, children, siblings, stacked_types,\n",
    "             mask_e, mask_d,\n",
    "             label_smooth = 1.0):\n",
    "        \n",
    "        output_enc, hn_enc = self.encode(input_word, input_char)\n",
    "        arc_c = tf.nn.elu(self.arc_c(output_enc))\n",
    "        type_c = tf.nn.elu(self.type_c(output_enc))\n",
    "        \n",
    "        output_dec, _ = self.decode(output_enc, heads, stacked_heads, siblings, hn_enc)\n",
    "        arc_h = tf.nn.elu(self.arc_h(output_dec))\n",
    "        type_h = tf.nn.elu(self.type_h(output_dec))\n",
    "        \n",
    "        max_len_d = tf.shape(arc_h)[1]\n",
    "        \n",
    "        out_arc = tf.squeeze(self.attention.forward(arc_h, arc_c, mask_d=mask_d, mask_e=mask_e), axis = 1)\n",
    "        batch = tf.shape(arc_c)[0]\n",
    "        max_len_e = tf.shape(arc_c)[1]\n",
    "        batch_index = tf.range(0, batch)\n",
    "        \n",
    "        t = tf.transpose(children)\n",
    "        broadcasted = tf.broadcast_to(batch_index, tf.shape(t))\n",
    "        concatenated = tf.transpose(tf.concat([tf.expand_dims(broadcasted, axis = 0), \n",
    "                                               tf.expand_dims(t, axis = 0)], axis = 0))\n",
    "        type_c = tf.gather_nd(type_c, concatenated)\n",
    "        out_type = self.bilinear.forward(type_h, type_c)\n",
    "        print(out_arc.shape,out_type.shape)\n",
    "        \n",
    "        minus_inf = -1e8\n",
    "        minus_mask_d = (1 - mask_d) * minus_inf\n",
    "        minus_mask_e = (1 - mask_e) * minus_inf\n",
    "        \n",
    "        out_arc = out_arc + tf.expand_dims(minus_mask_d, 2) + tf.expand_dims(minus_mask_e, 1)\n",
    "        loss_arc = tf.nn.log_softmax(out_arc, axis = 2)\n",
    "        loss_type = tf.nn.log_softmax(out_type, axis = 2)\n",
    "        coverage = tf.cumsum(tf.exp(loss_arc), axis = 1)\n",
    "        \n",
    "        mask_leaf = tf.cast(tf.equal(children, stacked_heads), tf.float32)\n",
    "        mask_non_leaf = (1.0 - mask_leaf)\n",
    "        \n",
    "        mask_d_2 = tf.expand_dims(mask_d, 2)\n",
    "        mask_e_1 = tf.expand_dims(mask_e, 1)\n",
    "        \n",
    "        loss_arc = loss_arc * mask_d_2 * mask_e_1\n",
    "        coverage = coverage * mask_d_2 * mask_e_1\n",
    "        loss_type = loss_type * mask_d_2\n",
    "        mask_leaf = mask_leaf * mask_d\n",
    "        mask_non_leaf = mask_non_leaf * mask_d\n",
    "        num_leaf = tf.reduce_sum(mask_leaf)\n",
    "        num_non_leaf = tf.reduce_sum(mask_non_leaf)\n",
    "        head_index = tf.tile(tf.expand_dims(tf.range(0, max_len_d), 1), [1, batch])\n",
    "        \n",
    "        t = tf.transpose(children)\n",
    "        broadcasted = tf.broadcast_to(batch_index, tf.shape(t))\n",
    "        concatenated = tf.transpose(tf.concat([tf.expand_dims(broadcasted, axis = 0),\n",
    "                                               tf.expand_dims(head_index, axis = 0),\n",
    "                                               tf.expand_dims(t, axis = 0)], axis = 0))\n",
    "        loss_arc = tf.gather_nd(loss_arc, concatenated)\n",
    "        \n",
    "        t = tf.transpose(stacked_types)\n",
    "        broadcasted = tf.broadcast_to(batch_index, tf.shape(t))\n",
    "        concatenated = tf.transpose(tf.concat([tf.expand_dims(broadcasted, axis = 0),\n",
    "                                               tf.expand_dims(head_index, axis = 0),\n",
    "                                               tf.expand_dims(t, axis = 0)], axis = 0))\n",
    "        loss_type = tf.gather_nd(loss_type, concatenated)\n",
    "        \n",
    "        loss_arc_leaf = loss_arc * mask_leaf\n",
    "        loss_arc_non_leaf = loss_arc * mask_non_leaf\n",
    "\n",
    "        loss_type_leaf = loss_type * mask_leaf\n",
    "        loss_type_non_leaf = loss_type * mask_non_leaf\n",
    "        \n",
    "        loss_cov = tf.clip_by_value(coverage - 2.0, 0.0, 100.0)\n",
    "        \n",
    "        return (tf.reduce_sum(-loss_arc_leaf) / num_leaf, \n",
    "                tf.reduce_sum(-loss_arc_non_leaf) / num_non_leaf,\n",
    "                tf.reduce_sum(-loss_type_leaf) / num_leaf, \n",
    "                tf.reduce_sum(-loss_type_non_leaf) / num_non_leaf,\n",
    "                tf.reduce_sum(loss_cov) / (num_leaf + num_non_leaf), \n",
    "                num_leaf, \n",
    "                num_non_leaf)\n",
    "    \n",
    "class Model:\n",
    "    def __init__(self, learning_rate = 1e-3, cov = 0.0):\n",
    "        self.stackpointer = StackPointer(word_dim = 128, \n",
    "                            num_words = len(word2idx), \n",
    "                            char_dim = 128, \n",
    "                            num_chars = len(char2idx), \n",
    "                            num_filters = 128, \n",
    "                            kernel_size = 3,\n",
    "                            input_size_decoder = 256, \n",
    "                            hidden_size = 256, \n",
    "                            layers = 1,\n",
    "                            num_labels = len(tag2idx), \n",
    "                            arc_space = 128, \n",
    "                            type_space = 128)\n",
    "        self.words = tf.placeholder(tf.int32, (None, None))\n",
    "        self.chars = tf.placeholder(tf.int32, (None, None, None))\n",
    "        self.heads = tf.placeholder(tf.int32, (None, None))\n",
    "        self.stacked_heads = tf.placeholder(tf.int32, (None, None))\n",
    "        self.siblings = tf.placeholder(tf.int32, (None, None))\n",
    "        self.childrens = tf.placeholder(tf.int32, (None, None))\n",
    "        self.stacked_types = tf.placeholder(tf.int32, (None, None))\n",
    "        self.mask_e = tf.placeholder(tf.float32, (None, None))\n",
    "        self.mask_d = tf.placeholder(tf.float32, (None, None))\n",
    "        loss_arc_leaf, loss_arc_non_leaf, \\\n",
    "        loss_type_leaf, loss_type_non_leaf, \\\n",
    "        loss_cov, num_leaf, num_non_leaf = self.stackpointer.loss(self.words, self.chars, self.heads, \n",
    "                               self.stacked_heads, self.childrens, \n",
    "                               self.siblings, self.stacked_types,\n",
    "                               self.mask_e, self.mask_d)\n",
    "        loss_arc = loss_arc_leaf + loss_arc_non_leaf\n",
    "        loss_type = loss_type_leaf + loss_type_non_leaf\n",
    "        self.cost = loss_arc + loss_type + cov * loss_cov\n",
    "        self.optimizer = tf.train.AdamOptimizer(learning_rate = learning_rate).minimize(self.cost)\n",
    "        \n",
    "        self.encode_output, self.encode_hidden = self.stackpointer.encode(self.words, self.chars)\n",
    "        self.encode_arc_c = tf.nn.elu(self.stackpointer.arc_c(self.encode_output))\n",
    "        self.type_c = tf.nn.elu(self.stackpointer.type_c(self.encode_output))\n",
    "        \n",
    "        self.src_encoding = tf.placeholder(tf.float32, (None, self.stackpointer.hidden_size))\n",
    "        self.arc_c = tf.placeholder(tf.float32, (None, self.stackpointer.arc_space))\n",
    "        self.hx = tf.placeholder(tf.float32, (None, \n",
    "                                              self.stackpointer.hidden_size * 2 * self.stackpointer.layers))      \n",
    "        \n",
    "        src_encoding = tf.nn.elu(self.stackpointer.src_dense(self.src_encoding))\n",
    "        output_dec, hx = self.stackpointer.decoder(src_encoding, self.hx)\n",
    "        arc_h = tf.nn.elu(self.stackpointer.arc_h(tf.expand_dims(output_dec, axis = 1)))\n",
    "        type_h = tf.nn.elu(self.stackpointer.type_h(output_dec))\n",
    "        out_arc = self.stackpointer.attention.forward(arc_h, tf.expand_dims(self.arc_c, 0))\n",
    "        out_arc = tf.squeeze(tf.squeeze(out_arc, axis = 1), axis = 1)\n",
    "        self.hyp_scores = tf.nn.log_softmax(out_arc, axis = 1)\n",
    "        self.type_h = type_h\n",
    "        self.decode_hidden = hx\n",
    "        \n",
    "        self.holder_type_h = tf.placeholder(tf.float32, (None, self.stackpointer.arc_space))\n",
    "        self.holder_type_c = tf.placeholder(tf.float32, (None, self.stackpointer.arc_space))\n",
    "        \n",
    "        out_type = self.stackpointer.bilinear.forward(self.holder_type_h, self.holder_type_c)\n",
    "        self.hyp_type_scores = tf.nn.log_softmax(out_type, axis = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-16-b0a50e6384bf>:73: LSTMCell.__init__ (from tensorflow.python.ops.rnn_cell_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "This class is equivalent as tf.keras.layers.LSTMCell, and will be replaced by that in Tensorflow 2.0.\n",
      "WARNING:tensorflow:<tensorflow.python.ops.rnn_cell_impl.LSTMCell object at 0x7f8991fc6b70>: Using a concatenated state is slower and will soon be deprecated.  Use state_is_tuple=True.\n",
      "WARNING:tensorflow:From <ipython-input-16-b0a50e6384bf>:83: MultiRNNCell.__init__ (from tensorflow.python.ops.rnn_cell_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "This class is equivalent as tf.keras.layers.StackedRNNCells, and will be replaced by that in Tensorflow 2.0.\n",
      "WARNING:tensorflow:<tensorflow.python.ops.rnn_cell_impl.LSTMCell object at 0x7f8991fe5278>: Using a concatenated state is slower and will soon be deprecated.  Use state_is_tuple=True.\n",
      "WARNING:tensorflow:<tensorflow.python.ops.rnn_cell_impl.LSTMCell object at 0x7f8991fe5240>: Using a concatenated state is slower and will soon be deprecated.  Use state_is_tuple=True.\n",
      "WARNING:tensorflow:From <ipython-input-16-b0a50e6384bf>:111: dynamic_rnn (from tensorflow.python.ops.rnn) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Please use `keras.layers.RNN(cell)`, which is equivalent to this API\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow/python/ops/rnn_cell_impl.py:961: calling Zeros.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Call initializer instance with the dtype argument instead of passing it to the constructor\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow/python/ops/init_ops.py:1251: calling VarianceScaling.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Call initializer instance with the dtype argument instead of passing it to the constructor\n",
      "(?, ?, ?) (?, ?, 52)\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow/python/ops/math_grad.py:1250: add_dispatch_support.<locals>.wrapper (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n"
     ]
    }
   ],
   "source": [
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "model = Model()\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_X = words_train\n",
    "train_Y = labels_train\n",
    "train_depends = depends_train\n",
    "train_char = sentences_train\n",
    "\n",
    "test_X = words_test\n",
    "test_Y = labels_test\n",
    "test_depends = depends_test\n",
    "test_char = sentences_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "prior_order = model.stackpointer.prior_order\n",
    "\n",
    "def decode_sentence(output_enc, arc_c, type_c, hx, beam, length, ordered, leading_symbolic):\n",
    "    def valid_hyp(base_id, child_id, head):\n",
    "        if constraints[base_id, child_id]:\n",
    "            return False\n",
    "        elif not ordered or prior_order == PriorOrder.DEPTH or child_orders[base_id, head] == 0:\n",
    "            return True\n",
    "        elif prior_order == PriorOrder.LEFT2RIGTH:\n",
    "            return child_id > child_orders[base_id, head]\n",
    "        else:\n",
    "            if child_id < head:\n",
    "                return child_id < child_orders[base_id, head] < head\n",
    "            else:\n",
    "                return child_id > child_orders[base_id, head]\n",
    "        \n",
    "    length = output_enc.shape[0] if length is None else length\n",
    "            \n",
    "    stacked_heads = [[0] for _ in range(beam)]\n",
    "    grand_parents = [[0] for _ in range(beam)]\n",
    "    siblings = [[0] for _ in range(beam)]\n",
    "    children = np.zeros((beam, 2 * length - 1))\n",
    "    stacked_types = np.zeros((beam, 2 * length - 1))\n",
    "    \n",
    "    children = np.zeros((beam, 2 * length - 1))\n",
    "    stacked_types = np.zeros((beam, 2 * length - 1))\n",
    "    hypothesis_scores = [0]\n",
    "    constraints = np.zeros([beam, length], dtype=np.bool)\n",
    "    constraints[:, 0] = True\n",
    "    child_orders = np.zeros([beam, length], dtype=np.int64)\n",
    "\n",
    "    new_stacked_heads = [[] for _ in range(beam)]\n",
    "    new_grand_parents = [[] for _ in range(beam)]\n",
    "    new_siblings = [[] for _ in range(beam)]\n",
    "    new_skip_connects = [[] for _ in range(beam)]\n",
    "    new_children = np.zeros((beam, 2 * length - 1))\n",
    "    new_stacked_types = np.zeros((beam, 2 * length - 1))\n",
    "    num_hyp = 1\n",
    "    num_step = 2 * length - 1\n",
    "    for t in range(num_step):\n",
    "        heads = np.array([stacked_heads[i][-1] for i in range(num_hyp)])\n",
    "        gpars = np.array([grand_parents[i][-1] for i in range(num_hyp)])\n",
    "        sibs = np.array([siblings[i].pop() for i in range(num_hyp)])\n",
    "        src_encoding = output_enc[heads]\n",
    "        mask_sibs = np.expand_dims((np.array(sibs) != 0).astype(np.float32), axis = 1)\n",
    "        output_enc_sibling = output_enc[sibs] * mask_sibs\n",
    "        src_encoding = src_encoding + output_enc_sibling\n",
    "        output_enc_gpar = output_enc[gpars]\n",
    "        src_encoding = src_encoding + output_enc_gpar\n",
    "        hyp_scores, type_h, hx = sess.run([model.hyp_scores, model.type_h, model.decode_hidden],\n",
    "                                     feed_dict = {model.src_encoding: src_encoding,\n",
    "                                                  model.arc_c: arc_c,\n",
    "                                                  model.hx: hx})\n",
    "        \n",
    "        new_hypothesis_scores = np.expand_dims(hypothesis_scores[:num_hyp], axis = 1) + hyp_scores\n",
    "        new_hypothesis_scores = new_hypothesis_scores.reshape((-1))\n",
    "        hyp_index = np.argsort(new_hypothesis_scores)[::-1]\n",
    "        new_hypothesis_scores = np.sort(new_hypothesis_scores)[::-1]\n",
    "        base_index = (hyp_index // length)\n",
    "        child_index = hyp_index % length\n",
    "        cc = 0\n",
    "        ids = []\n",
    "        new_constraints = np.zeros([beam, length], dtype=np.bool)\n",
    "        new_child_orders = np.zeros([beam, length], dtype=np.int64)\n",
    "        for id_ in range(num_hyp * length):\n",
    "            base_id = base_index[id_]\n",
    "            if base_id:\n",
    "                ids.append(id_)\n",
    "                continue\n",
    "            child_id = child_index[id_]\n",
    "            head = heads[base_id]\n",
    "            new_hyp_score = new_hypothesis_scores[id_]\n",
    "            if child_id == head:\n",
    "                if head != 0 or t + 1 == num_step:\n",
    "                    new_constraints[cc] = constraints[base_id]\n",
    "                    new_child_orders[cc] = child_orders[base_id]\n",
    "\n",
    "                    new_stacked_heads[cc] = [stacked_heads[base_id][i] for i in range(len(stacked_heads[base_id]))]\n",
    "                    new_stacked_heads[cc].pop()\n",
    "\n",
    "                    new_grand_parents[cc] = [grand_parents[base_id][i] for i in range(len(grand_parents[base_id]))]\n",
    "                    new_grand_parents[cc].pop()\n",
    "\n",
    "                    new_siblings[cc] = [siblings[base_id][i] for i in range(len(siblings[base_id]))]\n",
    "\n",
    "                    new_children[cc] = children[base_id]\n",
    "                    new_children[cc, t] = child_id\n",
    "\n",
    "                    hypothesis_scores[cc] = new_hyp_score\n",
    "                    ids.append(id_)\n",
    "                    cc += 1\n",
    "            elif valid_hyp(base_id, child_id, head):\n",
    "                new_constraints[cc] = constraints[base_id]\n",
    "                new_constraints[cc, child_id] = True\n",
    "\n",
    "                new_child_orders[cc] = child_orders[base_id]\n",
    "                new_child_orders[cc, head] = child_id\n",
    "\n",
    "                new_stacked_heads[cc] = [stacked_heads[base_id][i] for i in range(len(stacked_heads[base_id]))]\n",
    "                new_stacked_heads[cc].append(child_id)\n",
    "\n",
    "                new_grand_parents[cc] = [grand_parents[base_id][i] for i in range(len(grand_parents[base_id]))]\n",
    "                new_grand_parents[cc].append(head)\n",
    "\n",
    "                new_siblings[cc] = [siblings[base_id][i] for i in range(len(siblings[base_id]))]\n",
    "                new_siblings[cc].append(child_id)\n",
    "                new_siblings[cc].append(0)\n",
    "\n",
    "                new_children[cc] = children[base_id]\n",
    "                new_children[cc, t] = child_id\n",
    "\n",
    "                hypothesis_scores[cc] = new_hyp_score\n",
    "                ids.append(id_)\n",
    "                cc += 1\n",
    "                    \n",
    "            if cc == beam:\n",
    "                break\n",
    "            \n",
    "        num_hyp = len(ids)\n",
    "        if num_hyp == 0:\n",
    "            return None\n",
    "        else:\n",
    "            index = np.array(ids)\n",
    "        base_index = base_index[index]\n",
    "        child_index = child_index[index]\n",
    "        hyp_type_scores = sess.run(model.hyp_type_scores,\n",
    "                                  feed_dict = {\n",
    "                                      model.holder_type_h: type_h[base_index],\n",
    "                                      model.holder_type_c: type_c[child_index]\n",
    "                                  })\n",
    "        hyp_types = np.argmax(hyp_type_scores, axis = 1)\n",
    "        hyp_type_scores = np.max(hyp_type_scores, axis = 1)\n",
    "        hypothesis_scores[:num_hyp] = hypothesis_scores[:num_hyp] + hyp_type_scores\n",
    "\n",
    "        for i in range(num_hyp):\n",
    "            base_id = base_index[i]\n",
    "            new_stacked_types[i] = stacked_types[base_id]\n",
    "            new_stacked_types[i, t] = hyp_types[i]\n",
    "\n",
    "        stacked_heads = [[new_stacked_heads[i][j] for j in range(len(new_stacked_heads[i]))] for i in range(num_hyp)]\n",
    "        grand_parents = [[new_grand_parents[i][j] for j in range(len(new_grand_parents[i]))] for i in range(num_hyp)]\n",
    "        siblings = [[new_siblings[i][j] for j in range(len(new_siblings[i]))] for i in range(num_hyp)]\n",
    "        constraints = new_constraints\n",
    "        child_orders = new_child_orders\n",
    "        children = np.copy(new_children)\n",
    "        stacked_types = np.copy(new_stacked_types)\n",
    "        \n",
    "    children = children[0].astype(np.int32)\n",
    "    stacked_types = stacked_types[0].astype(np.int32)\n",
    "    heads = np.zeros(length, dtype=np.int32)\n",
    "    types = np.zeros(length, dtype=np.int32)\n",
    "    stack = [0]\n",
    "    for i in range(num_step):\n",
    "        head = stack[-1]\n",
    "        child = children[i]\n",
    "        type_ = stacked_types[i]\n",
    "        if child != head:\n",
    "            heads[child] = head\n",
    "            types[child] = type_\n",
    "            stack.append(child)\n",
    "        else:\n",
    "            stacked_types[i] = 0\n",
    "            stack.pop()\n",
    "\n",
    "    return heads, types, length, children, stacked_types   \n",
    "        \n",
    "def decode(input_word, input_char, length = None, beam = 1, leading_symbolic=0, ordered=True):\n",
    "    \n",
    "    arc_c, type_c, output, hn = sess.run([model.encode_arc_c, model.type_c, \n",
    "                                          model.encode_output, model.encode_hidden],\n",
    "            feed_dict = {model.words: input_word, model.chars: input_char})\n",
    "    batch, max_len_e, _ = output.shape\n",
    "\n",
    "    heads = np.zeros([batch, max_len_e], dtype=np.int32)\n",
    "    types = np.zeros([batch, max_len_e], dtype=np.int32)\n",
    "\n",
    "    children = np.zeros([batch, 2 * max_len_e - 1], dtype=np.int32)\n",
    "    stack_types = np.zeros([batch, 2 * max_len_e - 1], dtype=np.int32)\n",
    "    \n",
    "    for b in range(batch):\n",
    "        sent_len = None if length is None else length[b]\n",
    "        preds = decode_sentence(output[b], arc_c[b], type_c[b], [hn[b]], \n",
    "                                beam, sent_len, ordered, leading_symbolic)\n",
    "        if preds is None:\n",
    "            preds = decode_sentence(output[b], arc_c[b], type_c[b], [hn[b]], beam, \n",
    "                                         sent_len, False, leading_symbolic)\n",
    "        hids, tids, sent_len, chids, stids = preds\n",
    "        heads[b, :sent_len] = hids\n",
    "        types[b, :sent_len] = tids\n",
    "\n",
    "        children[b, :2 * sent_len - 1] = chids\n",
    "        stack_types[b, :2 * sent_len - 1] = stids\n",
    "\n",
    "    return heads, types, children, stack_types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_char_seq(batch, UNK = 2):\n",
    "    maxlen_c = max([len(k) for k in batch])\n",
    "    x = [[len(i) for i in k] for k in batch]\n",
    "    maxlen = max([j for i in x for j in i])\n",
    "    temp = np.zeros((len(batch),maxlen_c,maxlen),dtype=np.int32)\n",
    "    for i in range(len(batch)):\n",
    "        for k in range(len(batch[i])):\n",
    "            for no, c in enumerate(batch[i][k]):\n",
    "                temp[i,k,-1-no] = char2idx.get(c, UNK)\n",
    "    return temp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((5, 37), (5, 73))"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from tensorflow.keras.preprocessing.sequence import pad_sequences\n",
    "\n",
    "batch_x = train_X[:5]\n",
    "batch_x = pad_sequences(batch_x,padding='post')\n",
    "batch_char = train_char[:5]\n",
    "batch_char = generate_char_seq(batch_char)\n",
    "batch_y = train_Y[:5]\n",
    "batch_y = pad_sequences(batch_y,padding='post')\n",
    "batch_depends = train_depends[:5]\n",
    "batch_depends = pad_sequences(batch_depends,padding='post')\n",
    "\n",
    "batch_stacked_heads = stacked_heads_train[:5]\n",
    "batch_stacked_heads = pad_sequences(batch_stacked_heads,padding='post')\n",
    "batch_children = children_train[:5]\n",
    "batch_children = pad_sequences(batch_children,padding='post')\n",
    "batch_siblings = siblings_train[:5]\n",
    "batch_siblings = pad_sequences(batch_siblings,padding='post')\n",
    "batch_stacked_types = stacked_types_train[:5]\n",
    "batch_stacked_types = pad_sequences(batch_stacked_types,padding='post')\n",
    "batch_e = np.zeros(batch_x.shape)\n",
    "batch_d = np.zeros(batch_stacked_heads.shape)\n",
    "nonzero = np.count_nonzero(batch_x, axis = 1)\n",
    "\n",
    "for no, i in enumerate(nonzero):\n",
    "    batch_e[no,:i] = 1.0\n",
    "for no, i in enumerate(nonzero * 2 - 1):\n",
    "    batch_d[no,:i] = 1.0\n",
    "    \n",
    "batch_x.shape, batch_stacked_heads.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "14.264593"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "feed_dict = {model.words: batch_x,\n",
    "             model.chars: batch_char,\n",
    "             model.heads: batch_depends,\n",
    "             model.stacked_heads: batch_stacked_heads,\n",
    "             model.childrens: batch_children,\n",
    "             model.siblings: batch_siblings,\n",
    "             model.stacked_types: batch_stacked_types,\n",
    "             model.mask_e: batch_e,\n",
    "             model.mask_d: batch_d}\n",
    "sess.run(model.cost, feed_dict = feed_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.27 s, sys: 251 ms, total: 2.52 s\n",
      "Wall time: 1.32 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(array([[ 0,  0,  1,  0,  1,  6,  1,  1,  7,  0,  0, 12,  0,  0, 15,  8,\n",
       "         18, 18,  7, 21, 21, 18, 23, 21,  1, 28, 28, 28, 21,  1, 34, 34,\n",
       "         31, 34,  0, 34,  0],\n",
       "        [ 0, 10,  3, 10,  7,  7,  7,  3, 10, 10,  0, 10, 10, 14, 10, 16,\n",
       "         14, 10, 10, 23,  0,  0, 36, 36, 36, 36, 36, 36, 36, 36, 36, 36,\n",
       "         36, 36, 36, 36, 21],\n",
       "        [ 0,  0,  1,  4,  5,  1,  9,  9,  9,  5,  9, 13, 13,  9, 13, 16,\n",
       "         14,  1,  0,  0, 25,  0,  0, 25, 25, 22,  0,  0,  0,  0,  0,  0,\n",
       "          0,  0,  0,  0,  0],\n",
       "        [ 0,  6,  3,  1,  6,  6,  0,  9,  9,  6, 12, 12,  9, 15, 15, 12,\n",
       "          6,  0, 17,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "          0,  0,  0,  0,  0],\n",
       "        [ 0,  2,  6,  4,  2,  6,  0, 10, 10, 10,  6,  6, 16, 16, 16, 17,\n",
       "          6, 16, 17, 18, 18, 22, 17, 27, 27, 27, 27, 22, 31, 31, 31, 27,\n",
       "         35, 35, 35, 22,  6]], dtype=int32),\n",
       " array([[ 0,  5,  7,  8,  7, 13, 35, 28, 10,  8, 44,  7, 38,  7,  3, 35,\n",
       "          2,  3,  4,  2,  3, 14,  2, 14,  7,  2,  3, 13, 14,  7,  7,  7,\n",
       "          0,  7, 23, 23, 45],\n",
       "        [ 0,  7,  3,  6,  2,  3, 13, 14, 18, 18,  5, 10, 10,  2,  4, 11,\n",
       "         20,  7,  7,  0, 16,  4,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "          0,  0,  0,  0,  4],\n",
       "        [ 0,  5,  7, 13,  6, 28, 11,  6, 18, 27, 36,  9, 13, 10, 20,  2,\n",
       "          4,  7,  7, 23,  0, 23, 23,  0,  0,  4, 23, 23, 23, 23, 23, 23,\n",
       "         23, 23, 23, 23, 45],\n",
       "        [ 0,  6,  2, 14, 18, 19,  5,  2,  9,  4,  2,  3, 14,  2,  3, 14,\n",
       "          7, 28,  4,  4,  4,  4,  4,  4,  4,  4,  4,  4,  4,  4,  4,  4,\n",
       "          4, 23, 23, 23,  4],\n",
       "        [ 0,  3,  6,  2, 14, 25,  5,  2,  3, 15,  4,  7, 16,  6, 18,  0,\n",
       "         28, 12,  6,  8,  0, 16, 27,  2,  3, 13, 15, 10,  2,  3, 15, 14,\n",
       "          2,  2,  3,  4,  7]], dtype=int32),\n",
       " array([[ 1,  7, 18, 21, 23, 22, 22, 23, 28, 25, 25, 26, 26, 27, 27, 28,\n",
       "         19, 19, 20, 20, 21, 16, 16, 17, 17, 18,  8, 15, 14, 14, 15,  8,\n",
       "          7,  6,  5,  5,  6,  2,  2,  4,  4, 24, 24, 29, 29,  1, 36, 36,\n",
       "         34, 35, 35, 30, 30, 31, 32, 32, 31, 33, 33, 34,  3,  3, 10, 10,\n",
       "         12, 11, 11, 12,  9,  9, 13, 13,  0],\n",
       "        [10,  3,  7,  4,  4,  5,  5,  6,  6,  7,  2,  2,  3, 14, 16, 15,\n",
       "         15, 16, 13, 13, 14,  1,  1,  8,  8,  9,  9, 11, 11, 12, 12, 17,\n",
       "         17, 18, 18, 10, 20, 20, 21, 36, 22, 22, 23, 19, 19, 23, 24, 24,\n",
       "         25, 25, 26, 26, 27, 27, 28, 28, 29, 29, 30, 30, 31, 31, 32, 32,\n",
       "         33, 33, 34, 34, 35, 35, 36, 21,  0],\n",
       "        [ 1,  5,  9, 13, 14, 16, 15, 15, 16, 14, 11, 11, 12, 12, 13,  6,\n",
       "          6,  7,  7,  8,  8, 10, 10,  9,  4,  3,  3,  4,  5,  2,  2, 17,\n",
       "         17,  1, 36, 36, 35, 35, 34, 34, 33, 33, 32, 32, 31, 31, 30, 30,\n",
       "         29, 29, 28, 28, 27, 27, 26, 26, 21, 21, 22, 25, 23, 23, 24, 24,\n",
       "         20, 20, 25, 22, 19, 19, 18, 18,  0],\n",
       "        [ 6,  9, 12, 15, 13, 13, 14, 14, 15, 10, 10, 11, 11, 12,  7,  7,\n",
       "          8,  8,  9,  1,  3,  2,  2,  3,  1,  4,  4,  5,  5, 16, 16,  6,\n",
       "         36, 36, 35, 35, 34, 34, 33, 33, 32, 32, 31, 31, 30, 30, 29, 29,\n",
       "         28, 28, 27, 27, 26, 26, 25, 25, 24, 24, 23, 23, 22, 22, 21, 21,\n",
       "         20, 20, 19, 19, 17, 18, 18, 17,  0],\n",
       "        [ 6, 16, 17, 22, 27, 31, 28, 28, 29, 29, 30, 30, 31, 23, 23, 24,\n",
       "         24, 25, 25, 26, 26, 27, 35, 32, 32, 33, 33, 34, 34, 35, 21, 21,\n",
       "         22, 18, 19, 19, 20, 20, 18, 15, 15, 17, 12, 12, 13, 13, 14, 14,\n",
       "         16,  2,  4,  3,  3,  4,  1,  1,  2, 10,  7,  7,  8,  8,  9,  9,\n",
       "         10,  5,  5, 11, 11, 36, 36,  6,  0]], dtype=int32),\n",
       " array([[ 5, 28,  4, 14, 14,  2,  0,  0, 14,  2,  0,  3,  0, 13,  0,  0,\n",
       "          2,  0,  3,  0,  0,  2,  0,  3,  0,  0, 10, 35,  3,  0,  0,  0,\n",
       "          0, 35, 13,  0,  0,  7,  0,  7,  0,  7,  0,  7,  0,  0, 45,  0,\n",
       "         23, 23,  0,  7,  0,  7,  0,  0,  0,  7,  0,  0,  8,  0, 44,  0,\n",
       "         38,  7,  0,  0,  8,  0,  7,  0,  0],\n",
       "        [ 5,  6, 14,  2,  0,  3,  0, 13,  0,  0,  3,  0,  0,  4, 20, 11,\n",
       "          0,  0,  2,  0,  0,  7,  0, 18,  0, 18,  0, 10,  0, 10,  0,  7,\n",
       "          0,  7,  0,  0, 16,  0,  4,  4,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "          0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "        [ 5, 28, 27, 10, 20,  4,  2,  0,  0,  0,  9,  0, 13,  0,  0, 11,\n",
       "          0,  6,  0, 18,  0, 36,  0,  0,  6, 13,  0,  0,  0,  7,  0,  7,\n",
       "          0,  0, 45,  0, 23,  0, 23,  0, 23,  0, 23,  0, 23,  0, 23,  0,\n",
       "         23,  0, 23,  0, 23,  0, 23,  0, 23,  0, 23,  4,  0,  0,  0,  0,\n",
       "          0,  0,  0,  0, 23,  0,  7,  0,  0],\n",
       "        [ 5,  4, 14, 14,  2,  0,  3,  0,  0,  2,  0,  3,  0,  0,  2,  0,\n",
       "          9,  0,  0,  6, 14,  2,  0,  0,  0, 18,  0, 19,  0,  7,  0,  0,\n",
       "          4,  0, 23,  0, 23,  0, 23,  0,  4,  0,  4,  0,  4,  0,  4,  0,\n",
       "          4,  0,  4,  0,  4,  0,  4,  0,  4,  0,  4,  0,  4,  0,  4,  0,\n",
       "          4,  0,  4,  0, 28,  4,  0,  0,  0],\n",
       "        [ 5, 28, 12, 27, 10, 14,  2,  0,  3,  0, 15,  0,  0,  2,  0,  3,\n",
       "          0, 13,  0, 15,  0,  0,  4,  2,  0,  2,  0,  3,  0,  0, 16,  0,\n",
       "          0,  6,  8,  0,  0,  0,  0,  0,  0,  0, 16,  0,  6,  0, 18,  0,\n",
       "          0,  6, 14,  2,  0,  0,  3,  0,  0,  4,  2,  0,  3,  0, 15,  0,\n",
       "          0, 25,  0,  7,  0,  7,  0,  0,  0]], dtype=int32))"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "decode(batch_x, batch_char)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:01<00:00,  6.11it/s, cost=2.97]\n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 19.70it/s, cost=12.1]\n",
      "train minibatch loop:   0%|          | 1/375 [00:00<01:01,  6.08it/s, cost=3.28]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0, training loss: 5.157737, valid loss: 11.861909\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:00<00:00,  6.15it/s, cost=2.01]\n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 18.53it/s, cost=13.8]\n",
      "train minibatch loop:   0%|          | 1/375 [00:00<01:07,  5.52it/s, cost=2.31]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1, training loss: 2.576627, valid loss: 13.340673\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:00<00:00,  6.17it/s, cost=1.55] \n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 19.33it/s, cost=15.4]\n",
      "train minibatch loop:   0%|          | 1/375 [00:00<01:02,  5.99it/s, cost=1.77]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 2, training loss: 1.922838, valid loss: 14.725556\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:01<00:00,  6.11it/s, cost=1.36] \n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 19.20it/s, cost=16.4]\n",
      "train minibatch loop:   0%|          | 1/375 [00:00<01:05,  5.70it/s, cost=1.47]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 3, training loss: 1.529883, valid loss: 15.789502\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:00<00:00,  6.15it/s, cost=1.12] \n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 19.20it/s, cost=17.9]\n",
      "train minibatch loop:   0%|          | 1/375 [00:00<01:03,  5.88it/s, cost=1.2]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 4, training loss: 1.266019, valid loss: 17.307760\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:00<00:00,  6.16it/s, cost=1.02] \n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 18.87it/s, cost=19.5]\n",
      "train minibatch loop:   0%|          | 1/375 [00:00<01:03,  5.93it/s, cost=1.06]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 5, training loss: 1.066313, valid loss: 19.008535\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:00<00:00,  6.16it/s, cost=0.878]\n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 19.69it/s, cost=21.3]\n",
      "train minibatch loop:   0%|          | 1/375 [00:00<01:02,  6.03it/s, cost=0.895]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 6, training loss: 0.908035, valid loss: 20.994354\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:01<00:00,  6.13it/s, cost=0.748]\n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 18.86it/s, cost=22.6]\n",
      "train minibatch loop:   0%|          | 1/375 [00:00<01:02,  6.03it/s, cost=0.771]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 7, training loss: 0.780265, valid loss: 22.426714\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:01<00:00,  6.11it/s, cost=0.636]\n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 19.22it/s, cost=24.4]\n",
      "train minibatch loop:   0%|          | 1/375 [00:00<01:03,  5.93it/s, cost=0.615]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 8, training loss: 0.687402, valid loss: 24.419289\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:00<00:00,  6.17it/s, cost=0.628]\n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 18.42it/s, cost=26.8]\n",
      "train minibatch loop:   0%|          | 1/375 [00:00<01:05,  5.75it/s, cost=0.546]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 9, training loss: 0.609938, valid loss: 26.764641\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:00<00:00,  6.18it/s, cost=0.613]\n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 17.83it/s, cost=28.2]\n",
      "train minibatch loop:   0%|          | 1/375 [00:00<01:06,  5.65it/s, cost=0.52]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 10, training loss: 0.525183, valid loss: 28.478970\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:00<00:00,  6.15it/s, cost=0.538]\n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 18.67it/s, cost=31.2]\n",
      "train minibatch loop:   0%|          | 1/375 [00:00<01:06,  5.62it/s, cost=0.484]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 11, training loss: 0.459827, valid loss: 31.322876\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:00<00:00,  6.17it/s, cost=0.512] \n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 19.11it/s, cost=32.4]\n",
      "train minibatch loop:   0%|          | 1/375 [00:00<01:06,  5.59it/s, cost=0.367]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 12, training loss: 0.400364, valid loss: 33.366253\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:01<00:00,  6.14it/s, cost=0.413] \n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 18.76it/s, cost=34.1]\n",
      "train minibatch loop:   0%|          | 1/375 [00:00<01:02,  5.95it/s, cost=0.316]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 13, training loss: 0.357156, valid loss: 34.881569\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 375/375 [01:00<00:00,  6.16it/s, cost=0.331] \n",
      "test minibatch loop: 100%|██████████| 120/120 [00:06<00:00, 18.91it/s, cost=36.8]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 14, training loss: 0.307119, valid loss: 37.149876\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "batch_size = 32\n",
    "epoch = 15\n",
    "\n",
    "for e in range(epoch):\n",
    "    test_loss, train_loss = [], []\n",
    "    \n",
    "    pbar = tqdm(\n",
    "        range(0, len(train_X), batch_size), desc = 'train minibatch loop'\n",
    "    )\n",
    "    for i in pbar:\n",
    "        index = min(i + batch_size, len(train_X))\n",
    "        batch_x = train_X[i: index]\n",
    "        batch_x = pad_sequences(batch_x,padding='post')\n",
    "        batch_char = train_char[i: index]\n",
    "        batch_char = generate_char_seq(batch_char)\n",
    "        batch_y = train_Y[i: index]\n",
    "        batch_y = pad_sequences(batch_y,padding='post')\n",
    "        batch_depends = train_depends[i: index]\n",
    "        batch_depends = pad_sequences(batch_depends,padding='post')\n",
    "\n",
    "        batch_stacked_heads = stacked_heads_train[i: index]\n",
    "        batch_stacked_heads = pad_sequences(batch_stacked_heads,padding='post')\n",
    "        batch_children = children_train[i: index]\n",
    "        batch_children = pad_sequences(batch_children,padding='post')\n",
    "        batch_siblings = siblings_train[i: index]\n",
    "        batch_siblings = pad_sequences(batch_siblings,padding='post')\n",
    "        batch_stacked_types = stacked_types_train[i: index]\n",
    "        batch_stacked_types = pad_sequences(batch_stacked_types,padding='post')\n",
    "        batch_e = np.zeros(batch_x.shape)\n",
    "        batch_d = np.zeros(batch_stacked_heads.shape)\n",
    "        nonzero = np.count_nonzero(batch_x, axis = 1)\n",
    "\n",
    "        for no, i in enumerate(nonzero):\n",
    "            batch_e[no,:i] = 1.0\n",
    "        for no, i in enumerate(nonzero * 2 - 1):\n",
    "            batch_d[no,:i] = 1.0\n",
    "        \n",
    "        feed_dict = {model.words: batch_x,\n",
    "             model.chars: batch_char,\n",
    "             model.heads: batch_depends,\n",
    "             model.stacked_heads: batch_stacked_heads,\n",
    "             model.childrens: batch_children,\n",
    "             model.siblings: batch_siblings,\n",
    "             model.stacked_types: batch_stacked_types,\n",
    "             model.mask_e: batch_e,\n",
    "             model.mask_d: batch_d}\n",
    "        cost, _ = sess.run([model.cost, model.optimizer], feed_dict = feed_dict)\n",
    "        train_loss.append(cost)\n",
    "        pbar.set_postfix(cost = cost)\n",
    "        \n",
    "    pbar = tqdm(\n",
    "        range(0, len(test_X), batch_size), desc = 'test minibatch loop'\n",
    "    )\n",
    "    for i in pbar:\n",
    "        index = min(i + batch_size, len(test_X))\n",
    "        batch_x = test_X[i: index]\n",
    "        batch_x = pad_sequences(batch_x,padding='post')\n",
    "        batch_char = test_char[i: index]\n",
    "        batch_char = generate_char_seq(batch_char)\n",
    "        batch_y = test_Y[i: index]\n",
    "        batch_y = pad_sequences(batch_y,padding='post')\n",
    "        batch_depends = test_depends[i: index]\n",
    "        batch_depends = pad_sequences(batch_depends,padding='post')\n",
    "\n",
    "        batch_stacked_heads = stacked_heads_test[i: index]\n",
    "        batch_stacked_heads = pad_sequences(batch_stacked_heads,padding='post')\n",
    "        batch_children = children_test[i: index]\n",
    "        batch_children = pad_sequences(batch_children,padding='post')\n",
    "        batch_siblings = siblings_test[i: index]\n",
    "        batch_siblings = pad_sequences(batch_siblings,padding='post')\n",
    "        batch_stacked_types = stacked_types_test[i: index]\n",
    "        batch_stacked_types = pad_sequences(batch_stacked_types,padding='post')\n",
    "        batch_e = np.zeros(batch_x.shape)\n",
    "        batch_d = np.zeros(batch_stacked_heads.shape)\n",
    "        nonzero = np.count_nonzero(batch_x, axis = 1)\n",
    "\n",
    "        for no, i in enumerate(nonzero):\n",
    "            batch_e[no,:i] = 1.0\n",
    "        for no, i in enumerate(nonzero * 2 - 1):\n",
    "            batch_d[no,:i] = 1.0\n",
    "        \n",
    "        feed_dict = {model.words: batch_x,\n",
    "             model.chars: batch_char,\n",
    "             model.heads: batch_depends,\n",
    "             model.stacked_heads: batch_stacked_heads,\n",
    "             model.childrens: batch_children,\n",
    "             model.siblings: batch_siblings,\n",
    "             model.stacked_types: batch_stacked_types,\n",
    "             model.mask_e: batch_e,\n",
    "             model.mask_d: batch_d}\n",
    "        cost = sess.run(model.cost, feed_dict = feed_dict)\n",
    "        test_loss.append(cost)\n",
    "        pbar.set_postfix(cost = cost)\n",
    "    \n",
    "    print(\n",
    "    'epoch: %d, training loss: %f, valid loss: %f\\n'\n",
    "    % (e, np.mean(train_loss), np.mean(test_loss)))\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(heads_pred, types_pred, heads, types, lengths,\n",
    "             symbolic_root=False, symbolic_end=False):\n",
    "    batch_size, _ = heads_pred.shape\n",
    "    ucorr = 0.\n",
    "    lcorr = 0.\n",
    "    total = 0.\n",
    "    ucomplete_match = 0.\n",
    "    lcomplete_match = 0.\n",
    "\n",
    "    corr_root = 0.\n",
    "    total_root = 0.\n",
    "    start = 1 if symbolic_root else 0\n",
    "    end = 1 if symbolic_end else 0\n",
    "    for i in range(batch_size):\n",
    "        ucm = 1.\n",
    "        lcm = 1.\n",
    "        for j in range(start, lengths[i] - end):\n",
    "\n",
    "            total += 1\n",
    "            if heads[i, j] == heads_pred[i, j]:\n",
    "                ucorr += 1\n",
    "                if types[i, j] == types_pred[i, j]:\n",
    "                    lcorr += 1\n",
    "                else:\n",
    "                    lcm = 0\n",
    "            else:\n",
    "                ucm = 0\n",
    "                lcm = 0\n",
    "\n",
    "            if heads[i, j] == 0:\n",
    "                total_root += 1\n",
    "                corr_root += 1 if heads_pred[i, j] == 0 else 0\n",
    "\n",
    "        ucomplete_match += ucm\n",
    "        lcomplete_match += lcm\n",
    "    \n",
    "    return ucorr / total, lcorr / total, corr_root / total_root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.6045627376425855, 0.5209125475285171, 0.90625)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "heads, types, _, _ = decode(batch_x, batch_char)\n",
    "arc_accuracy, type_accuracy, root_accuracy = evaluate(heads, types, batch_depends, batch_y, \n",
    "                                                      np.count_nonzero(batch_x, axis = 1))\n",
    "arc_accuracy, type_accuracy, root_accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "arcs, types, roots = [], [], []\n",
    "\n",
    "for i in range(0, len(test_X), 5):\n",
    "    index = min(i + 5, len(test_X))\n",
    "    batch_x = test_X[i: index]\n",
    "    batch_x = pad_sequences(batch_x,padding='post')\n",
    "    batch_char = test_char[i: index]\n",
    "    batch_char = generate_char_seq(batch_char)\n",
    "    batch_y = test_Y[i: index]\n",
    "    batch_y = pad_sequences(batch_y,padding='post')\n",
    "    batch_depends = test_depends[i: index]\n",
    "    batch_depends = pad_sequences(batch_depends,padding='post')\n",
    "    \n",
    "    heads, tags_seq, _, _ = decode(batch_x, batch_char)\n",
    "    \n",
    "    arc_accuracy, type_accuracy, root_accuracy = evaluate(heads, tags_seq, batch_depends, batch_y, \n",
    "                                                      np.count_nonzero(batch_x, axis = 1))\n",
    "    arcs.append(arc_accuracy)\n",
    "    types.append(type_accuracy)\n",
    "    roots.append(root_accuracy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "arc accuracy: 0.6188156085110088\n",
      "types accuracy: 0.482035002661857\n",
      "root accuracy: 0.8939869281045753\n"
     ]
    }
   ],
   "source": [
    "print('arc accuracy:', np.mean(arcs))\n",
    "print('types accuracy:', np.mean(types))\n",
    "print('root accuracy:', np.mean(roots))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
